import os
import sys
from copy import deepcopy
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.api as sm
from sklearn.metrics import mean_absolute_error
from dateutil.relativedelta import relativedelta
import warnings
warnings.filterwarnings('ignore') # 常用
from numpy.linalg import LinAlgError

import functools
import time

join = os.path.join
dirname = os.path.dirname

CURRENT_PATH = dirname(os.path.realpath(__file__))
DATA_PATH = join(dirname(CURRENT_PATH) , 'DataAssets')

'''
使用SARIMAX算法建模，使用 经销商配件销售量 遍历寻找最优参数。
1. 外生性（虚拟）变量设定为 异常点，以虚拟变量0-1输入。（可配置不使用）
2. 不使用周期性分解

input :pd.series 时序数据
output : trainedmodel

'''


def time_log(text):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kw):
            start_time = time.time()
            time_str = time.strftime(
                u"%Y-%m-%d %X", time.localtime(start_time))
            print(' %s: %s Starting!' % (time_str, text))
            ret = func(*args, **kw)
            end_time = time.time()
            time_str = time.strftime(
                u"%Y-%m-%d %X", time.localtime(end_time))
            print(' %s: %s Finish!\n' % (time_str, text))
            print(' 耗时: %s 秒!\n' % (end_time - start_time))
            return ret
        return wrapper
    return decorator


class Arima_Model:
    def __init__(self, endog, \
                 model_name, \
                 exog=None, \
                 out_of_var=None, \
                 seasonal_order=None,
                 maxLag=3, \
                 forecast_lag=12):

        """
        :param endog: pd.Series 预测的变量
        :param model_name: 模型的名称
        :param exog:  目前只使用 0-1 虚拟变量代表事件
        :param out_of_var: pd.Series   目前只使用 0-1 虚拟变量代表事件，外推预测时使用，未提供默认元素值均为0
        :param maxLag: p，q的最大搜索范围
        :param predict_lag: 向前预测的间隔 ，
        """

        self.data_ts = endog
        self.model_name = model_name

        self.pred_error = None
        self.pred_mae = None
        self.pred_ci = None
        self.pred_ts = None

        self.forecast_ts = None
        self.forecast_ci = None

        self.forecast_lag = forecast_lag
        self.maxLag = maxLag
        self.p = 1
        self.q = 1
        self.d = 1
        self.properModel = {}
        self.bic = sys.maxsize
        self.efor_inverti = True
        self.efor_station = True
        self.error_rate = None
        self.time_diff = 1

        self.out_of_var = out_of_var

        if out_of_var is not None and  len(out_of_var) != self.forecast_lag:
            raise('len(out_of_var) != len(forecast_lag)')
        elif out_of_var is None:
            index = self.data_ts.index[-1]
            index_list = [ index + relativedelta(months=i) for i in range(1 , self.forecast_lag + 1)]
            out_of_var = pd.Series([0] * self.forecast_lag , index = index_list)
            self.out_of_var = out_of_var
        else:
            index = self.data_ts.index[-1]
            index_list = [ index + relativedelta(months=i) for i in range(1 , self.forecast_lag + 1)]
            self.out_of_var = out_of_var
            self.out_of_var.index = index_list

        if exog is not None:
            if len(exog) == len(self.data_ts):
                self.exog = exog
            else:
                raise('len(enog) != len(exog)')
        else:
            self.exog = pd.DataFrame({'exog': 0}, index=endog.index)

        if seasonal_order is None:
            self.seasonal_order = (0, 0, 0, 0)
        elif self.data_ts.shape[0] > 12:
            if self.data_ts.shape[0] > 24:
                self.seasonal_order = (1, self.d, 0, 12)
            else:
                self.seasonal_order = (1, self.d, 0, 4)


    # 历史值的对比
    def get_train_predict(self):
        '''
        非预测，调用模型的预测方法获取现有序列的预测结果
        '''
        start_id = 0
        ex_ts_final = sm.add_constant(self.exog, has_constant='add')

        predict = self.properModel[self.model_name].get_prediction(exog=ex_ts_final)

        # 时间序列预测误差值
        predict_error = predict.predicted_mean.iloc[start_id:] - self.data_ts.iloc[start_id:]

        self.pred_ts = deepcopy(predict.predicted_mean[start_id:])
        self.pred_error = deepcopy(predict_error)
        self.pred_ci = deepcopy(predict.conf_int()[start_id:])
        self.pred_mae = mean_absolute_error(self.data_ts.iloc[start_id:], predict.predicted_mean.iloc[start_id:] )
        print('model MAE: ',self.pred_mae)

    # 对未来指定间隔的预测
    def get_forcest(self):
        """
        调用模型进行预测，未知时间点的预测
        :return:
        """
        start_id = self.data_ts.shape[0]
        end_id = self.data_ts.shape[0] + self.forecast_lag - 1

        # 自变量没有赋值，需要先对自变量进行预测
        if self.out_of_var is None:
            out_of_full = sm.add_constant(pd.Series([0] * (self.forecast_lag)), has_constant='add')
        else:
            out_of_full = sm.add_constant(self.out_of_var, has_constant='add')

        forecast = self.properModel[self.model_name].get_prediction(start=start_id, end=end_id, exog=out_of_full)
        self.forecast_ts = deepcopy(forecast.predicted_mean)
        self.forecast_ci = deepcopy(forecast.conf_int())

    # 绘制训练结果和训练预测值的对比
    def draw_result(self, if_retrain):
        if if_retrain:
            self.get_train_predict()
        fig, ax = plt.subplots(figsize=(12, 8))
        # 绘制点数据
        self.data_ts.plot(ax=ax, style='.', label='Observed')

        # Plot predictions
        self.pred_ts.plot(ax=ax, style='r--', label='trian_data')
        ci = self.pred_ci
        ax.fill_between(ci.index, ci.iloc[:, 0], ci.iloc[:, 1], color='r', alpha=0.05)

        # 预测值预测
        self.forecast_ts.plot(ax=ax, style='g--', label='predict_data')
        ci = self.forecast_ci
        ax.fill_between(ci.index, ci.iloc[:, 0], ci.iloc[:, 1], color='g', alpha=0.05)

        ax.legend(loc='lower right')

    @time_log("模型训练")
    def gen_model_auto(self):
        ex_ts_final = sm.add_constant(self.exog, has_constant='add')

        for d in np.arange(2):
            for p in np.arange(self.maxLag):
                for q in np.arange(self.maxLag):
                    if p > 0 and q > 0:
                        self.efor_inverti = True
                        self.efor_station = True
                        try:
                            # enforce_stationarity=True, enforce_invertibility=True 有可能会导致矩阵分解失败从而报错
                            # 一方面是statsmodel的bug，另一方面设定的滞后项越多越难用少量数据估计这么多参数
                            # exog 中提供了常数项，模型中不在定义
                            mod = sm.tsa.statespace.SARIMAX(self.data_ts, \
                                                            exog=ex_ts_final, \
                                                            order=(p, d, q), \
                                                            seasonal_order=self.seasonal_order,
                                                            enforce_stationarity = self.efor_station,
                                                            enforce_invertibility = self.efor_inverti)
                            results_ARMA = mod.fit(disp=False)
                        except LinAlgError:
                            self.efor_inverti = False
                            self.efor_station = False
                            mod = sm.tsa.statespace.SARIMAX(self.data_ts, \
                                                            exog=ex_ts_final, \
                                                            order=(p, d, q), \
                                                            seasonal_order=self.seasonal_order,
                                                            enforce_stationarity = self.efor_station,
                                                            enforce_invertibility = self.efor_inverti)
                            results_ARMA = mod.fit(disp=False)
                        bic = results_ARMA.bic

                        if bic < self.bic:
                            self.p = p
                            self.q = q
                            self.d = d
                            self.bic = bic


        mod = sm.tsa.statespace.SARIMAX(self.data_ts, \
                                        exog=ex_ts_final, \
                                        order=(self.p, self.d, self.q), \
                                        seasonal_order=self.seasonal_order,
                                        enforce_stationarity=self.efor_station,
                                        enforce_invertibility=self.efor_inverti
                                        )


        self.properModel = {self.model_name: results_ARMA}


    def gen_model_manual(self, parm_dict):

        p = parm_dict['p']
        d = parm_dict['d']
        q = parm_dict['q']
        seasonal_order = parm_dict['seasonal_order']

        if seasonal_order is None:
            seasonal_order = (0, 0, 0, 0)

        ex_ts = sm.add_constant(self.exog, has_constant='add')

        try:
            mod = sm.tsa.statespace.SARIMAX(self.data_ts, exog = ex_ts, order=(p, d, q),
                                            seasonal_order=seasonal_order)
            results_ARMA = mod.fit(disp=False)
        except Exception:
            mod = sm.tsa.statespace.SARIMAX(self.data_ts, exog = ex_ts, order=(p, d, q),
                                            seasonal_order=seasonal_order, \
                                            enforce_stationarity=False, \
                                            enforce_invertibility=False)
            results_ARMA = mod.fit(disp=False)

        bic = results_ARMA.bic

        self.p = p
        self.q = q
        self.d = d
        self.properModel = {'original': results_ARMA}
        self.bic = bic


if __name__ == '__main__':

    cur_name = 'test_arima'

    cur_file = pd.read_excel(join(DATA_PATH, 'cur_ts_melt.xlsx'))
    cur_file = cur_file.sort_values(by = ['asc_code' , 'date'] , ascending= True)
    cur_file.index = cur_file['date']

    cur_file_one = cur_file.loc[cur_file['asc_code'] == 0000000,'value']


    out_of_var = pd.Series([0, 1, 0])  # 这里应该是预测之后12月的销量，考虑到春节事件
    arima_model_one = Arima_Model(endog = cur_file_one,
                                   model_name = 'model_test',\
                                   exog = None,
                                   out_of_var = out_of_var,\
                                   seasonal_order = None,
                                   maxLag = 3,
                                   forecast_lag = 3)

    arima_model_one.gen_model_auto()
    arima_model_one.get_train_predict() # 历史数据的预测

    arima_model_one.draw_result(False) # 作图查看结果


